from manim import *

class GeometricSeriesSum(Scene):
    def construct(self):
        tex_template = TexTemplate()
        tex_template.add_to_preamble(r"\usepackage{ctex}")
        
        # 等比数列的定义和公式
        title = Text("等比数列求和", font_size=48)
        formula = MathTex(
            r"S_n = a + ar + ar^2 + \cdots + ar^{n-1} = a \frac{1-r^n}{1-r}",
            tex_template=tex_template
        ).scale(0.8)

        
        # 展示标题和公式
        self.play(Write(title))
        self.wait(1)
        self.play(Transform(title, formula))
        self.wait(2)
        # title向上移动
        self.play(title.animate.shift(UP))

        title2 = Text("令a=1简化公式", font_size=30).next_to(title, DOWN)
        formula2 = MathTex(
            r"S_n = 1 + r + r^2 + \cdots + r^{n-1} = \frac{r^n - 1}{r - 1}",
            tex_template=tex_template
        ).scale(0.8).next_to(title2, DOWN)
         # 公式也可以写成r*Sn-Sn = 1+r+r^2+r^3+r^4 + r^{n-1}
        formula3 = MathTex(
            r"rS_n - S_n = r^{n}-1",
            tex_template=tex_template
        ).scale(0.8).next_to(formula2, DOWN)

        self.play(FadeIn(title2))
        self.wait(2)
        self.play(FadeIn(formula2))
        self.wait(2)
        self.play(FadeIn(formula3))
        self.wait(2)
        

        # 等比数列的例子
        example=Text("演示公比为2，项数为4的求和过程如下", font_size=30).next_to(formula3, DOWN)
        self.play(FadeIn(example))
        self.wait(2)

        self.play(FadeOut(title),FadeOut(formula2),FadeOut(formula3),FadeOut(example),FadeOut(title2))

        # 设置等比数列的第一个元素和公比
        a1 = 1
        r = 2
        n = 4  # 项数
        # 计算最大值以动态调整单位长度，确保坐标轴在屏幕内
        max_value = 15
        unit_length = 6.5 / max_value  # 3 是一个经验值，可以根据需要调整
        # 创建坐标轴，扩大范围
        axes = Axes(
            x_range=[0, n + 2, 1],  # x 轴范围稍微扩大一些
            y_range=[0, max_value + 2, 1],  # y 轴范围稍微扩大一些
            x_length=(n + 2) * unit_length,
            y_length=(max_value + 2) * unit_length,
            axis_config={"include_numbers": True},
        ).add_coordinates().shift(UP * 0.5).scale(0.8)
        axes_labels = axes.get_axis_labels(x_label="n", y_label="a_n")
        self.play(Create(axes), Write(axes_labels))

        # 创建等比数列的正方形表示，把正方形放在坐标轴上
        squares = VGroup()
        for i in range(n):
            num_squares = a1 * (r ** i)
            group = VGroup(*[Square(side_length=unit_length*0.8) for _ in range(num_squares)])
            group.arrange(DOWN, buff=0).next_to(axes.c2p(i + 0.5, 0), UP, buff=0)
            squares.add(group)
        
        # 动画展示正方形
        self.play(Create(squares))
        self.wait(2)
        
        # 显示问题，如果把数列所有项都乘以公比，会怎么样呢?
        question = Text("所有项都乘以公比r").scale(0.8).next_to(axes, LEFT, buff=0.5)
        self.play(Write(question))
        self.wait(1)   

        # 复制坐标轴和所有的正方形，放在原来的坐标轴右侧 
        axes_copy = axes.copy().shift(RIGHT * (n + 3) * unit_length)
        squares_copy = squares.copy().shift(RIGHT * (n + 3) * unit_length)
        
        self.play(Create(axes_copy), Create(squares_copy))
        self.wait(2)

        # 在新坐标轴上将正方形数量依次变为原来的 r 倍
        for i in range(n):
            num_squares = a1 * (r ** (i + 1))
            new_group = VGroup(*[Square(side_length=unit_length*0.8) for _ in range(num_squares)])
            new_group.arrange(DOWN, buff=0).next_to(axes_copy.c2p(i + 0.5, 0), UP, buff=0)
            self.play(Transform(squares_copy[i], new_group))
            self.wait(1)
            
        # 观察2个数列重叠的部分
        question1 = Text("观察重叠的部分").scale(0.8).next_to(question, DOWN, buff=0.5)
        self.play(Write(question1))
        self.wait(1)           

        # 突出显示第一个坐标轴上的 2, 4, 8 对应的正方形
        highlight_indices = [1, 2, 3]  # 对应的索引为 1, 2, 3
        highlight_color = YELLOW

        for i in highlight_indices:
            self.play(squares[i].animate.set_fill(highlight_color, opacity=0.8))
            self.play(squares_copy[i-1].animate.set_fill(highlight_color, opacity=0.8))
            self.wait(1)

        # 在第一个坐标轴下方使用brace显示公式 Sn，在第二个坐标轴下使用brace显示r*Sn    
        brace1 = Brace(squares, DOWN)
        text1 = brace1.get_text("$S_n$")
        self.play(GrowFromCenter(brace1), Write(text1))
        self.wait(1)

        # 在第二个坐标轴下方使用brace显示公式 r*Sn
        brace2 = Brace(squares_copy, DOWN)
        text2 = brace2.get_text("$r \\cdot S_n$")
        self.play(GrowFromCenter(brace2), Write(text2))
        self.wait(1)
        
        # 在第一个坐标轴左侧显示公式：r*Sn-Sn=？
        self.play(FadeOut(question),FadeOut(question1))
        formula4 = MathTex(
            r"r \cdot S_n - S_n = ?",
            tex_template=tex_template
        ).scale(0.8)
        formula4.next_to(axes, LEFT, buff=0.5)
        self.play(Write(formula4))
        self.wait(1)    
        
        self.play(squares_copy[3].animate.set_fill(RED, opacity=0.8))
        self.wait(1) 
        self.play(squares[0].animate.set_fill(GREEN, opacity=0.8))
        self.wait(1) 
        formula3.next_to(axes, LEFT, buff=0.5)
        self.play(Transform(formula4,formula3))
        self.wait(1)  
        formula5 = MathTex(
            r"S_n = \frac{r^n - 1}{r - 1}",
            tex_template=tex_template
        ).scale(0.8).next_to(formula4, DOWN)
        self.play(Write(formula5))
        self.wait(2)
        
        
        
          